"""
Praat音频处理模块
使用parselmouth库提取音频基频
"""
import parselmouth
import numpy as np
from typing import List, Tuple, Optional
import os

class PraatProcessor:
    def __init__(self, f0_min: float = 75, f0_max: float = 500, time_step: float = 0.01):
        """
        初始化Praat处理器
        
        Args:
            f0_min: 最小基频 (Hz)
            f0_max: 最大基频 (Hz) 
            time_step: 时间步长 (秒)
        """
        self.f0_min = f0_min
        self.f0_max = f0_max
        self.time_step = time_step
    
    def extract_f0(self, audio_path: str) -> Tuple[np.ndarray, np.ndarray]:
        """
        提取音频文件的基频序列
        
        Args:
            audio_path: 音频文件路径
            
        Returns:
            (time_array, f0_array): 时间序列和基频序列
        """
        try:
            # 读取音频文件
            sound = parselmouth.Sound(audio_path)
            
            # 提取基频
            pitch = sound.to_pitch(
                time_step=self.time_step,
                pitch_floor=self.f0_min,
                pitch_ceiling=self.f0_max
            )
            
            # 获取基频值
            f0_values = pitch.selected_array['frequency']
            time_values = pitch.xs()
            
            # 移除未定义的基频值 (NaN)
            valid_indices = ~np.isnan(f0_values)
            time_clean = time_values[valid_indices]
            f0_clean = f0_values[valid_indices]
            
            return time_clean, f0_clean
            
        except Exception as e:
            print(f"提取基频时发生错误: {e}")
            return np.array([]), np.array([])
    
    def smooth_f0(self, f0_array: np.ndarray, window_size: int = 5) -> np.ndarray:
        """
        平滑基频序列
        
        Args:
            f0_array: 基频序列
            window_size: 滑动窗口大小
            
        Returns:
            平滑后的基频序列
        """
        if len(f0_array) < window_size:
            return f0_array
            
        # 使用移动平均平滑
        smoothed = np.convolve(f0_array, np.ones(window_size)/window_size, mode='same')
        return smoothed
    
    def normalize_f0(self, f0_array: np.ndarray) -> np.ndarray:
        """
        标准化基频序列 (Z-score标准化)
        
        Args:
            f0_array: 基频序列
            
        Returns:
            标准化后的基频序列
        """
        if len(f0_array) == 0:
            return f0_array
            
        mean_f0 = np.mean(f0_array)
        std_f0 = np.std(f0_array)
        
        if std_f0 == 0:
            return f0_array - mean_f0
            
        return (f0_array - mean_f0) / std_f0
    
    def extract_f0_contour(self, audio_path: str, normalize: bool = False, smooth: bool = True) -> Tuple[np.ndarray, np.ndarray]:
        """
        提取完整的基频轮廓
        
        Args:
            audio_path: 音频文件路径
            normalize: 是否进行标准化
            smooth: 是否进行平滑处理
            
        Returns:
            (time_array, f0_array): 处理后的时间序列和基频序列
        """
        time_array, f0_array = self.extract_f0(audio_path)
        
        if len(f0_array) == 0:
            return time_array, f0_array
        
        # 平滑处理
        if smooth:
            f0_array = self.smooth_f0(f0_array)
        
        # 标准化处理
        if normalize:
            f0_array = self.normalize_f0(f0_array)
        
        return time_array, f0_array
    
    def get_audio_info(self, audio_path: str) -> dict:
        """
        获取音频文件基本信息
        
        Args:
            audio_path: 音频文件路径
            
        Returns:
            包含音频信息的字典
        """
        try:
            sound = parselmouth.Sound(audio_path)
            
            return {
                'duration': sound.duration,
                'sampling_frequency': sound.sampling_frequency,
                'n_channels': sound.n_channels,
                'n_samples': sound.n_samples,
                'file_size': os.path.getsize(audio_path)
            }
        except Exception as e:
            print(f"获取音频信息时发生错误: {e}")
            return {}
    
    def validate_audio(self, audio_path: str) -> bool:
        """
        验证音频文件是否有效
        
        Args:
            audio_path: 音频文件路径
            
        Returns:
            音频文件是否有效
        """
        try:
            if not os.path.exists(audio_path):
                return False
                
            sound = parselmouth.Sound(audio_path)
            
            # 检查基本条件
            if sound.duration <= 0:
                return False
            if sound.sampling_frequency <= 0:
                return False
                
            # 尝试提取基频
            time_array, f0_array = self.extract_f0(audio_path)
            
            # 检查是否成功提取到基频
            return len(f0_array) > 0
            
        except Exception as e:
            print(f"验证音频文件时发生错误: {e}")
            return False

# 使用示例
if __name__ == "__main__":
    processor = PraatProcessor()
    
    # 测试音频文件路径（需要实际存在的文件）
    test_audio = "test.wav"
    
    if os.path.exists(test_audio):
        print("音频信息:")
        info = processor.get_audio_info(test_audio)
        for key, value in info.items():
            print(f"  {key}: {value}")
        
        print("\n提取基频:")
        time_array, f0_array = processor.extract_f0_contour(test_audio)
        print(f"  时间点数: {len(time_array)}")
        print(f"  基频点数: {len(f0_array)}")
        if len(f0_array) > 0:
            print(f"  基频范围: {np.min(f0_array):.2f} - {np.max(f0_array):.2f} Hz")
    else:
        print(f"测试音频文件不存在: {test_audio}")
